from src.Q1_特征提取 import data_process, feature_extract
from src.Q1_分类模型 import data_scale
from sklearn.ensemble import AdaBoostClassifier
import pandas as pd
import warnings
warnings.filterwarnings('ignore')


material_list = ['材料1', '材料2', '材料3', '材料4']
predict = []

for i in range(len(material_list)):
    # 训练模型
    df = pd.read_excel('../data/Q1_data.xlsx', sheet_name=material_list[i])
    X_train = data_scale(df)
    y_train = df['励磁波形']
    ada_clf = AdaBoostClassifier(n_estimators=50, random_state=42)
    ada_clf.fit(X_train, y_train)

    # 预测
    df = pd.read_excel('../2024年中国研究生数学建模竞赛赛题/C题/附件二（测试集）.xlsx')
    extracted_data = feature_extract(df)
    df = extracted_data
    X_test = data_scale(df[df['磁芯材料'] == material_list[i]])
    y_pred = ada_clf.predict(X_test)
    predict.extend(y_pred)

result_data = pd.read_excel('../2024年中国研究生数学建模竞赛赛题/C题/附件四（Excel表）.xlsx')
result_data.iloc[: 80, 1] = predict

counts = result_data.iloc[: 80, 1].value_counts().sort_index()
print(f'分类结果中，每种数量的个数为：\n{counts}')
# 保存结果
result_data.to_excel('../data/附件四（Excel表）.xlsx', index=False)

